import sys, os
import time
import string

from lxml import etree
import importlib, sys

# 设置默认字符集为UTF8 不然有些时候转码会出问题
default_encoding = 'utf-8'
if sys.getdefaultencoding() != default_encoding:
    importlib.reload(sys)
    sys.setdefaultencoding(default_encoding)


class XmlMgr:
    def __init__(self, root_node='root'):
        self.root = etree.Element(root_node)
        self.heads = []

    def add_element(self, name, value, head=None):
        if head == None:
            site = etree.SubElement(self.root, name)
        else:
            site = etree.SubElement(head, name)
        site.text = value

    def add_head(self, head_name, head=None):
        site = None
        if head == None:
            site = etree.SubElement(self.root, head_name)
        else:
            site = etree.SubElement(head, head_name)
        self.heads.append(site)
        return site

    def write(self, out_dir, out_name, verbose=False):
        dataxml = etree.tostring(self.root, pretty_print=True, encoding="UTF-8", method="xml", xml_declaration=True,
                                 standalone=None)
        if verbose:
            print(dataxml)
        if os.path.exists(out_dir) != True:
            os.makedirs(out_dir)
        f = open(os.path.join(out_dir, out_name), 'w')
        f.write(dataxml.decode('utf-8'))

    def gen_ann(self, filename, objects, width, height, depth):
        # objuects: (box.xmin, box.ymin,box.xmax, box.ymax,label_str)
        xml_gen = XmlMgr(root_node='annotation')

        xml_gen.add_element("filename", filename)
        size_head = xml_gen.add_head("size")
        xml_gen.add_element("width", str(width), size_head)
        xml_gen.add_element("height", str(height), size_head)
        xml_gen.add_element("depth", str(depth), size_head)

        for i in range(len(objects)):
            obj_head = xml_gen.add_head("object")
            xml_gen.add_element("name", str(objects[i][4]), obj_head)
            xml_gen.add_element("pose", "Unspecified", obj_head)
            xml_gen.add_element("truncated", "0", obj_head)
            xml_gen.add_element("difficult", "0", obj_head)
            box_head = xml_gen.add_head("bndbox", head=obj_head)
            xml_gen.add_element("xmin", str(objects[i][0]), box_head)
            xml_gen.add_element("ymin", str(objects[i][1]), box_head)
            xml_gen.add_element("xmax", str(objects[i][2]), box_head)
            xml_gen.add_element("ymax", str(objects[i][3]), box_head)
        xml_gen.write(os.path.dirname(filename), os.path.splitext(os.path.basename(filename))[0] + ".xml")


def create_xml():
    out_dir = 'out'
    data = etree.Element("data")
    # 1 interface_version
    interface_version_txt = '5'
    interface_version = etree.SubElement(data, 'interface_version')
    interface_version.text = interface_version_txt
    # 2 site
    site_txt = 'www.xxx.com'
    site = etree.SubElement(data, 'site')
    site.text = site_txt
    # 3 lastmod
    lastmod_txt = time.strftime('%Y-%m-%d', time.localtime())
    lastmod = etree.SubElement(data, 'lastmod')
    lastmod.text = lastmod_txt
    # 5 app
    app = etree.SubElement(data, 'app')
    # 6 title
    title_txt = u'%s' % '真心话大冒险'
    title_txt = etree.CDATA(title_txt)
    title = etree.SubElement(app, 'title')
    title.text = title_txt
    # 7 appid
    appid = etree.SubElement(app, 'appid')
    appid.text = '%s' % '222'

    dataxml = etree.tostring(data, pretty_print=True, encoding="UTF-8", method="xml", xml_declaration=True,
                             standalone=None)
    print(dataxml)
    if os.path.exists(out_dir) != True:
        os.makedirs(out_dir)
    f = open(os.path.join(out_dir, 'out.xml'), 'w')
    f.write(dataxml.decode('utf-8'))


def case1():
    xml_gen = XmlMgr(root_node='annotation')

    xml_gen.add_element("filename", "ffff")
    size_head = xml_gen.add_head("size")
    xml_gen.add_element("width", "223", size_head)
    xml_gen.add_element("height", "2223", size_head)
    xml_gen.add_element("depth", "3", size_head)

    for i in range(3):
        obj_head = xml_gen.add_head("object")
        xml_gen.add_element("name", "car", obj_head)
        xml_gen.add_element("pose", "Unspecified", obj_head)
        xml_gen.add_element("truncated", "0", obj_head)
        xml_gen.add_element("difficult", "0", obj_head)
        box_head = xml_gen.add_head("bndbox", head=obj_head)
        xml_gen.add_element("xmin", "632", box_head)
        xml_gen.add_element("ymin", "94", box_head)
        xml_gen.add_element("xmax", "771", box_head)
        xml_gen.add_element("ymax", "164", box_head)
    xml_gen.write("./gen", "test.xml")

# if __name__ == '__main__':
#     # create_xml()
#     case1()
